package com.flsldl.mydb.backend.vm;

import com.flsldl.mydb.common.Error;

import java.util.*;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * @ClassName LockTable
 * @Author @O_o @Date 2023/3/4 下午 7:11
 */
/*
* 维护一个依赖等待图，以进行死锁检测
* */
public class LockTable {

    // 某个 XID 已经获得的资源的 UID 列表
    private Map<Long, List<Long>> x2u;
    // UID 被某个 XID 持有
    private Map<Long, Long> u2x;
    // 正在等待 UID 的 XID 列表
    private Map<Long, List<Long>> wait;
    // 正在等待资源的 XID 的锁
    private Map<Long, Lock> waitLock;
    // XID 正在等待的 UID
    private Map<Long, Long> waitU;
    private Lock lock;

    public LockTable() {
        x2u = new HashMap<>();
        u2x = new HashMap<>();
        wait = new HashMap<>();
        waitLock = new HashMap<>();
        waitU = new HashMap<>();
        lock = new ReentrantLock();
    }

    // 不需要等待则返回 null， 否则返回锁对象
    // 会造成死锁则抛出异常
    public Lock add(long xid, long uid) throws Exception {
        lock.lock();
        try {
            // 1. 判断当前 XID 需要的 UID 资源是否已获取
            if(isInList(x2u, xid, uid)){
                return null;
            }
            // 2. 如果没有获取，判断当前 UID 是否被其他 XID 占用
            // fasle ==> 占用该 UID
            // containsKey 判断当前 map 集合里是否有该键 key
            if (!u2x.containsKey(uid)) {
                u2x.put(uid, xid);
                // 更新 x2u 集合, UID 资源已获取
                putIntoList(x2u, xid, uid);
                return null;
            }
            // 3. 如果已被其他 XID 占用， 进入等待
            waitU.put(xid, uid);
            putIntoList(wait, uid, xid);
            // 4. 死锁判断 true ==> 回滚，抛出异常
            if (hasDeadLock()) {
                waitU.remove(xid);
                removeFromList(wait, uid, xid);
                throw Error.DeadlockException;
            }
            // 5. 加锁，并等待
            Lock l = new ReentrantLock();
            l.lock();
            waitLock.put(xid, l);
            return l;
        } finally {
            lock.unlock();
        }
    }

    /**
     * 删除 XID 相关内容
     * @param xid
     */
    public void remove(long xid) {
        lock.lock();
        try {
            // 获取当前 XID 占用的 UID 资源列表
            List<Long> l = x2u.get(xid);
            // 删除
            if(l != null) {
                while(l.size() > 0) {
                    Long uid = l.remove(0);
                    selectNewXID(uid);
                }
            }
            waitU.remove(xid);
            x2u.remove(xid);
            waitLock.remove(xid);

        } finally {
            lock.unlock();
        }
    }

    // 从等待队列中选择一个xid来占用uid
    private void selectNewXID(long uid) {
        u2x.remove(uid);
        List<Long> l = wait.get(uid);
        if(l == null) return;
        assert l.size() > 0;

        while(l.size() > 0) {
            long xid = l.remove(0);
            if(!waitLock.containsKey(xid)) {
                continue;
            } else {
                u2x.put(uid, xid);
                Lock lo = waitLock.remove(xid);
                waitU.remove(xid);
                lo.unlock();
                break;
            }
        }

        if(l.size() == 0) wait.remove(uid);
    }

    // ----------------------------------(0v0)--> 分割线--------------------------------------------------------------

    private Map<Long, Integer> xidStamp;
    // 标记
    private int stamp;

    private boolean hasDeadLock() {

        xidStamp = new HashMap<>();
        stamp = 1;
        for (long xid : x2u.keySet()) {
            // 查询已获得资源的 XID 的锁标记
            Integer s = xidStamp.get(xid);
            // 锁标记 > 0 ==> 有锁且合法
            if (s != null && s > 0) {
                continue;
            }
            // 死锁判断
            stamp ++;
            if (dfs(xid)) {
                return true;
            }
        }
        return false;
    }

    /**
     * 死锁判断
     * @param xid
     * @return
     */
    private boolean dfs(long xid) {
        Integer stp = xidStamp.get(xid);
        // 死锁标记判断
        // 获得的锁标记重复 ==> 死锁
        if(stp != null && stp == stamp) {
            return true;
        }
        // 获得的锁标记合法 ==> ！死锁
        if(stp != null && stp < stamp) {
            return false;
        }
        // 当前 XID 锁标记为 null，添加锁标记
        xidStamp.put(xid, stamp);
        Long uid = waitU.get(xid);
        // 当前 XID 没有等待的资源 ==> !死锁
        if(uid == null) return false;
        // 等待的资源 UID 正在被哪个 XID 占用
        Long x = u2x.get(uid);
        // 如果该 XID不存在 报错
        assert x != null;
        // 循环判断该 XID 是否死锁
        return dfs(x);
    }

    // 将集合中下标为 uid0 位置上的 uid1 删除
    private void removeFromList(Map<Long, List<Long>> listMap, long uid0, long uid1) {
        List<Long> l = listMap.get(uid0);
        if(l == null) return;
        Iterator<Long> i = l.iterator();
        while(i.hasNext()) {
            long e = i.next();
            if(e == uid1) {
                i.remove();
                break;
            }
        }
        if(l.size() == 0) {
            listMap.remove(uid0);
        }
    }


    // ----------------------------------(0v0)--> 分割线--------------------------------------------------------------

    // 将 uid1 放入到集合中 下标为 uid0 的位置
    private void putIntoList(Map<Long, List<Long>> listMap, long uid0, long uid1) {

        if (!listMap.containsKey(uid0)) {
            listMap.put(uid0, new ArrayList<>());
        }
        listMap.get(uid0).add(0, uid1);
    }

    // 判断集合里是否存在 uid1
    private boolean isInList(Map<Long, List<Long>> listMap, long uid0, long uid1) {

        List<Long> l = listMap.get(uid0);
        if (l == null) return false;
        Iterator<Long> i = l.listIterator();
        while (i.hasNext()){
            long e= i.next();
            if (e == uid1) {
                return true;
            }
        }
        return false;

    }
}
